file(GLOB SOURCES "*.cpp")
set(PYTHON_EXECUTABLE "python3")
execute_process(
    COMMAND ${PYTHON_EXECUTABLE} -c "import sys; print([s for s in sys.path if 'packages' in s][0])"
    OUTPUT_VARIABLE PYTHON_SYS_PATH
)
string(STRIP ${PYTHON_SYS_PATH} PYTHON_SYS_PATH)

if("${PYTHON_SYS_PATH}" STREQUAL "")
    message(FATAL_ERROR "Python path is empty! Please check the python env.")
endif()

if (WITH_STORE)
    include_directories("../mooncake-store/include")
    include_directories("../mooncake-store/include/cachelib_memory_allocator")

    include_directories("../mooncake-store/include/cachelib_memory_allocator/include")
    include_directories("../mooncake-store/include/cachelib_memory_allocator/fake_include")
endif()

include_directories("/usr/include/jsoncpp")
include_directories("./")


set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE)
set(CMAKE_BUILD_WITH_INSTALL_RPATH TRUE)

message("${PYTHON_SYS_PATH}")

set(PYTHON_PACKAGE_NAME "mooncake")

pybind11_add_module(engine ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
    transfer_engine/transfer_engine_py.cpp 
)

target_link_libraries(engine PUBLIC 
    transfer_engine 
    glog::glog
    gflags::gflags
)

set(ALLOCATOR_SO_PATH "${CMAKE_BINARY_DIR}/mooncake-transfer-engine/nvlink-allocator/nvlink_allocator.so")
if(USE_MNNVL)
    message(STATUS "USE_MNNVL is enabled, nvlink_allocator.so will be installed in the Python package")
    install(FILES
        "${ALLOCATOR_SO_PATH}"
        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
    )
endif()

if (WITH_STORE)
    pybind11_add_module(store ${SOURCES} ${CACHE_ALLOCATOR_SOURCES}
        store/store_py.cpp
        integration_utils.h
    )
    target_link_libraries(store PUBLIC
        transfer_engine
        glog::glog
        gflags::gflags
        mooncake_store
        cachelib_memory_allocator
    )
endif()

message("${PYTHON_SYS_PATH}")

file(WRITE ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/__init__.py
    "# Auto-generated by CMake\n"
)

if (USE_MNNVL)
    message(STATUS "USE_MNNVL is enabled, allocator.py will be installed in the Python package")
    install(FILES
        "${CMAKE_CURRENT_SOURCE_DIR}/allocator.py"
        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
    )
endif()

# Install Python scripts from mooncake-wheel/mooncake/ directory
install(FILES
    "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/http_metadata_server.py"
    "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli_bench.py"
    "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/transfer_engine_topology_dump.py"
    "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_connector_v1.py"
    "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/vllm_v1_proxy_server.py"
    DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
)

if (WITH_STORE)
    install(FILES
        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_store_service.py"
        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/mooncake_config.py"
        "${CMAKE_CURRENT_SOURCE_DIR}/../mooncake-wheel/mooncake/cli.py"
        DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
    )
endif()

install(
    DIRECTORY
        ${CMAKE_CURRENT_BINARY_DIR}/${PYTHON_PACKAGE_NAME}/
    DESTINATION
        ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}
)

install(
    CODE "
        execute_process(COMMAND chmod 767 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}\")
        execute_process(COMMAND chmod 766 \"${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME}/__init__.py\")
    "
)

if (WITH_STORE)
    install(TARGETS store DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
endif()

install(TARGETS engine DESTINATION ${PYTHON_SYS_PATH}/${PYTHON_PACKAGE_NAME})
